package com.to8to.kitt.zk;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.concurrent.ConcurrentHashMap;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.KeeperException;
import org.apache.zookeeper.WatchedEvent;
import org.apache.zookeeper.Watcher;
import org.apache.zookeeper.ZooDefs.Ids;
import org.apache.zookeeper.ZooKeeper;
import org.apache.zookeeper.ZooKeeper.States;
import org.apache.zookeeper.data.Stat;

@SuppressWarnings({ "rawtypes", "unchecked", "unused" })
public class ZkClient implements Watcher {

	private static final Log logger = LogFactory.getLog("BaseLog");

	public static int ZK_NULL = 0;
	public static int ZK_CONNECTING = 1;
	public static int ZK_CONNECTED = 2;
	public static int ZK_DISCONNECTED = 3;

	protected ZooKeeper zookeeper;
	protected String hosts;
	protected int zkTimeout = 10000;
	protected int maxWaitSec = -1;

	private Map<String, String> ephemeralNodes = new ConcurrentHashMap<String, String>();
	private Map<String, ZkLock> lockMap = new ConcurrentHashMap<String, ZkLock>();
	private Map<String, byte[]> watchNodes = new ConcurrentHashMap<String, byte[]>(); // /< path---->data
	private Map<String, List<String>> watchDirectories = new ConcurrentHashMap<String, List<String>>(); // /< direcoty--->chind list

	ZkEventHandler bizhandler;

	private int state = ZK_NULL;

	public ZkClient() {
	}

	public ZkClient(ZkEventHandler handler) {
		this.bizhandler = handler;
	}

	public ZkEventHandler getBizhandler() {
		return bizhandler;
	}

	public void setBizhandler(ZkEventHandler bizhandler) {
		this.bizhandler = bizhandler;
	}

	public void setHosts(String hosts) {
		this.hosts = hosts;
	}

	public void setMaxWaitSec(int time) {
		this.maxWaitSec = time;
	}

	public void setZkTimeout(int time) {
		this.zkTimeout = time;
	}

	public void registEphemeralNode(String path, byte[] data) {
		logger.info("registEphemeralNode, path:" + path);

		ephemeralNodes.put(path, data != null ? new String(data) : "");
	}

	public void registLockNode(String fold, String lockPrefix) {
		logger.info("registLockNode: " + fold);

		lockMap.put(fold, new ZkLock(fold, lockPrefix));
	}

	public void registWatchNode(String path) {
		logger.info("registWatchNode: " + path);

		watchNodes.put(path, null);
	}

	public void unregistWatchNode(String path) {
		logger.info("unregistWatchNode: " + path);
		watchNodes.remove(path);
	}

	public void registWatchDirectory(String path) {
		logger.info("registWatchDirectory: " + path);
		List<String> list = new ArrayList<String>();
		watchDirectories.put(path, list);
	}

	public void unregistWatchDiretory(String path) {
		logger.info("unregistWatchDiretory: " + path);
		watchDirectories.remove(path);
	}

	public boolean initialize() {
		logger.info("init start!");
		try {
			logger.info("connectint to zk: " + hosts + " ...... \n");
			zookeeper = new ZooKeeper(hosts, zkTimeout, this);
			state = ZK_CONNECTING;
			int waitTime = 0;
			while (zookeeper.getState() != States.CONNECTED) {
				Thread.sleep(1000);
				++waitTime;
				if (maxWaitSec > 0 && waitTime == maxWaitSec)
					break;
			}
			if (zookeeper.getState() == States.CONNECTED) {
				state = ZK_CONNECTED;
				logger.info("zk connected!");
			} else {
				logger.error("zk connect failed!");
				return false;
			}
		} catch (Exception e) {
			logger.error(e);
			return false;
		}

		createEpemeralNodes();

		initializeWatchers();

		logger.info("init success!");

		return true;
	}

	public int getState() {
		return state;
	}

	public void process(WatchedEvent event) {

		logger.debug("<<<<" + event.toString() + ">>>>");

		if (event.getType() == Event.EventType.None) {
			logger.trace("Event type: EventType.None");
			switch (event.getState()) {
			case SyncConnected: {
				logger.trace("Event state: SyncConnected");
				break;
			}
			case Expired: {
				logger.trace("Event state: Expired");
				break;
			}
			case Disconnected: {
				logger.trace("Event state: Disconnected");
				if (state != ZK_CONNECTING) {
					state = ZK_DISCONNECTED;
					initialize();
				}
			}
			}
		} else {
			String path = event.getPath();
			switch (event.getType()) {
			case NodeDataChanged: {
				logger.trace("Event type: NodeDataChanged");
				watchData(path);
				break;
			}
			case NodeChildrenChanged: {
				logger.trace("Event type: NodeChildrenChanged");
				watchDir(path);
				break;
			}
			case NodeCreated: {
				logger.trace("Event type: NodeChildrenChanged");
				break;
			}
			case NodeDeleted: {
				logger.trace("Event type: NodeDeleted");
				break;
			}
			}
		}
	}

	protected void createEpemeralNodes() {
		logger.info("createEpemeralNodes...start!");

		for (Entry<String, String> entry : ephemeralNodes.entrySet()) {
			String path = entry.getKey();
			try {
				if (!exists(path, 10)) {
					String ret = zookeeper.create(path, entry.getValue().getBytes(), Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL);
					logger.info("create ephemeral node:[" + ret + "]");
				} else
					logger.warn("node exists:" + path);

			} catch (Exception e) {
				logger.error(e);
			}
		}

		// /< create lock nodes
		for (ZkLock lock : lockMap.values()) {
			boolean create = false;
			String ret = null;
			try {
				logger.debug("process local old lock:" + lock.getFullPath());
				if (lock.isNull()) {
					logger.debug("|----->lock is null!");
					create = true;
				} else if (!exists(lock.getFullPath(), 10)) {
					logger.debug("|----->lock not exist!");
					create = true;
				}
				if (create) {
					ret = zookeeper.create(lock.getFullPrefix(), null, Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL_SEQUENTIAL);
					logger.debug("create ephemeral sequential node: [" + ret + "]");
				} else
					logger.warn("|----->lock already exists:" + lock.getFullPath());
			} catch (Exception e) {
				logger.error(e);
			}
			if (create) {
				if (ret != null && ret.length() > 0) {
					String lockId = parseLockId(ret);
					lock.setLockId(lockId);
					logger.debug("set lock id:" + lockId);
				} else
					logger.error("no ret for create sequential node:" + lock.getFullPrefix());
			}
		}
		logger.info("createEpemeralNodes ... finish!");
	}

	protected void initializeWatchers() {
		logger.info("initializeWatchers...start!");
		for (String path : watchNodes.keySet())
			watchData(path);
		for (String path : watchDirectories.keySet())
			watchDir(path);
		logger.info("initializeWatchers..finish!");
	}

	protected boolean checkConnection() {
		logger.trace("check zk connection ... ");

		if (state == ZK_DISCONNECTED || state == ZK_NULL) {
			logger.warn("zookeeper is null or zookeeper is disconnected, try to initialize it!");
			return initialize();
		} else {
			int waitTime = 0;
			while (state == ZK_CONNECTING) {
				try {
					Thread.sleep(1000);
					logger.trace("...");
					++waitTime;
					if (maxWaitSec > 0 && waitTime == maxWaitSec)
						break;
				} catch (Exception e) {
					logger.error(e);
				}
			}
			if (state == ZK_CONNECTED) {
				logger.trace("zk conntected!");
				return true;
			} else {
				logger.error("zk connection failed!");
				return false;
			}
		}
	}

	protected byte[] watchData(String path) {
		logger.trace("watch node data, path = " + path);

		byte[] data = null;
		try {
			Stat stat = zookeeper.exists(path, false);
			if (stat == null) {
				logger.error("watch node not exist:" + path);
				return null;
			}
			data = zookeeper.getData(path, true, stat);
		} catch (Exception e) {
			logger.error(e);
		}

		if (bizhandler != null) {
			byte[] oldData = this.watchNodes.get(path);

			if ((data == null && oldData != null) || (oldData == null && data != null) || !data.equals(oldData)) {
				logger.trace("data changed!");
				bizhandler.onDataChanged(path, data);
			}
		}

		this.watchNodes.put(path, data);
		return data;
	}

	protected List<String> watchDir(String path) {
		logger.trace("watch directory, path = " + path);
		List<String> childrenNew = null;
		try {
			childrenNew = zookeeper.getChildren(path, true);
		} catch (Exception e) {
			logger.error(e);
		}

		if (bizhandler != null && childrenNew != null) {
			List<String> childrenOld = this.watchDirectories.get(path);

			if (childrenOld == null)
				childrenOld = new ArrayList<String>();

			List<String> added = new ArrayList<String>();
			List<String> removed = new ArrayList<String>();
			for (String tmp : childrenNew) {
				if (!childrenOld.contains(tmp))
					added.add(tmp);
			}
			for (String tmp : childrenOld) {
				if (!childrenNew.contains(tmp))
					removed.add(tmp);
			}
			if (added.size() > 0)
				bizhandler.onChildAdded(path, added);
			if (removed.size() > 0)
				bizhandler.onChildRemoved(path, removed);
		}

		watchDirectories.put(path, childrenNew);

		return childrenNew;

	}

	public boolean createNode(String path, byte[] data, CreateMode mode) {
		logger.info("try to create node: " + path);

		if (!checkConnection()) {
			logger.error("zk connection failed!");
			return false;
		}

		try {
			Stat stat = zookeeper.exists(path, null);
			if (stat != null) {
				logger.warn("node already exists: " + path);
				return true;
			}
			String ret = zookeeper.create(path, data, Ids.OPEN_ACL_UNSAFE, mode);
			logger.info("node created: " + path);
			return ret.equals(path);
		} catch (Exception e) {
			logger.error(e);
			return false;
		}
	}

	public boolean createDir(String path) {
		logger.info("try to create directory: " + path);

		if (!checkConnection()) {
			logger.error("zk connection failed!");
			return false;
		}

		if (path.length() <= 1 || path.charAt(0) != '/') {
			logger.error("invalid path:" + path);
			return false;
		}
		if (!path.endsWith("/"))
			path = path + "/";
		String dir = "/";
		boolean ret = true;
		for (int i = 1; i < path.length(); dir += path.charAt(i), ++i) {
			if (path.charAt(i) == '/') {
				ret = createNode(dir, null, CreateMode.PERSISTENT);
				if (!ret) {
					logger.error("create node failed : " + dir);
					break;
				}
			}
		}
		return ret;
	}

	public boolean removeDir(String path) {
		logger.info("try to remove directory: " + path);
		if (!checkConnection()) {
			logger.error("zk connection failed!");
			return false;
		}

		return removeDirRecursively(path);
	}

	protected boolean removeDirRecursively(String path) {
		logger.info("remove directory or node recursively: " + path);

		try {
			Stat stat = zookeeper.exists(path, null);
			if (stat == null) {
				logger.warn("remove failed, path not exists:" + path);
				return false;
			}
			ArrayList<String> children = (ArrayList<String>) zookeeper.getChildren(path, false);
			if (children != null && children.size() > 0) {
				String prefix = path.equals("/") ? "" : path;
				for (String child : children) {
					if (!removeDirRecursively(prefix + "/" + child))
						return false;
				}
				children = (ArrayList<String>) zookeeper.getChildren(path, false);
				if (children == null || children.size() == 0)
					zookeeper.delete(path, -1);
			} else
				zookeeper.delete(path, -1);
		} catch (Exception e) {
			logger.error(e);
			return false;
		}
		return true;
	}

	public List<String> listChildren(String path) {
		if (!checkConnection())
			return null;

		try {
			return (ArrayList<String>) zookeeper.getChildren(path, false);
		} catch (Exception e) {
			logger.error(e);
			return null;
		}
	}

	public boolean setData(String path, String data) {
		logger.info("set data, path: " + path + ", data:" + data);

		if (!checkConnection())
			return false;

		try {
			Stat stat = zookeeper.setData(path, data.getBytes(), -1);
			return true;
		} catch (Exception e) {
			logger.error(e);
		}
		return false;
	}

	public boolean exists(String path) {
		return exists(path, 5);
	}

	public boolean exists(String path, int repeatSec) {
		if (!checkConnection())
			return false;

		try {
			for (int i = 0; i < repeatSec; ++i) {
				Stat stat = zookeeper.exists(path, false);
				if (stat == null)
					return false;
				Thread.sleep(1000);
			}
		} catch (Exception e) {
			logger.error(e);
		}
		return true;
	}

	public String getData(String path) {
		if (!checkConnection())
			return null;

		Stat stat = new Stat();
		byte[] data = null;
		try {
			data = zookeeper.getData(path, false, stat);
		} catch (KeeperException e) {
			logger.error("KeeperException!", e);
		} catch (InterruptedException e) {
			logger.error("InterruptedException!", e);
		}
		if (data != null)
			return new String(data);
		return null;
	}

	public boolean printDir(String path, int dep) {
		String tab = "";
		for (int i = 0; i < dep; ++i)
			tab += "  ";

		int index = path.lastIndexOf("/");

		String node = null;
		if (index != -1)
			node = path.substring(index);
		else
			node = path;

		logger.info(tab + node);

		Stat stat = null;
		try {
			stat = zookeeper.exists(path, null);
		} catch (KeeperException e) {
			logger.error("KeeperException!", e);
		} catch (InterruptedException e) {
			logger.error("InterruptedException!", e);
		}
		if (stat == null) {
			logger.warn("path not exists:" + path);
			return false;
		}

		ArrayList<String> children = null;
		try {
			children = (ArrayList<String>) zookeeper.getChildren(path, false);
		} catch (KeeperException e) {
			logger.error("KeeperException!", e);
		} catch (InterruptedException e) {
			logger.error("InterruptedException!", e);
		}
		if (children != null && children.size() > 0) {
			String prefix = path.equals("/") ? "" : path;
			for (String child : children) {
				if (!printDir(prefix + "/" + child, dep + 1))
					return false;
			}
		}
		return true;
	}

	public boolean acquireLock(String path) {
		if (!this.checkConnection())
			return false;

		ZkLock lock = lockMap.get(path);

		if (lock == null || lock.isNull()) {
			logger.error("not registed lock yet, please regist it first!");
			return false;
		}

		String myname = lock.getLocalName();

		while (true) {
			try {
				List<String> children = zookeeper.getChildren(path, false);
				if (children == null || children.size() == 0) {
					logger.error("no children error! please regist it first!");
					return false;
				}

				Collections.sort(children, new Comparator() {
					public int compare(Object a, Object b) {
						int orderA = Integer.valueOf(parseLockId((String) a));
						int orderB = Integer.valueOf(parseLockId((String) b));
						return orderA - orderB;
					}
				});

				int myindex = children.indexOf(myname);

				if (myindex == 0) {
					logger.info("ok, I am the earliest node now! lock successfully acquired!");
					return true;
				}

				String wpath = path + "/" + children.get(myindex - 1);
				Stat stat = zookeeper.exists(wpath, true);
				logger.info("watch node data:" + wpath);
				this.registWatchNode(wpath);

				if (stat == null) {
					logger.info("previous node [" + path + "/" + children.get(myindex - 1) + "] not exists, try to reacquire...");
				} else {
//					logger.info("wait for previous node to quit: " + path + "/" + children.get(myindex - 1));
					Thread.sleep(5000);
				}
			} catch (KeeperException e) {
				logger.error("KeeperException!", e);
			} catch (InterruptedException e) {
				logger.error("InterruptedException!", e);
			}
		}
	}

	private String parseLockId(String path) {
		int index = path.lastIndexOf("-");
		if (index != -1)
			return path.substring(index + 1);
		return "";
	}

	/**
	 * @param args
	 * @throws InterruptedException
	 * @throws KeeperException
	 * @throws IOException
	 */
	public static void main(String[] args) throws IOException, KeeperException, InterruptedException {

		ZkClient zk = new ZkClient();

		zk.setHosts("192.168.20.23:9090");

		zk.registLockNode("/testlock", "lock");

		while (true) {
			logger.info("try to acquire lock: /testlock ...");
			if (!zk.acquireLock("/testlock"))
				logger.error("acquire lock error, check!");
			logger.info("acquire the lock, go!");
			Thread.sleep(5000);
		}
	}
}
